#!/bin/bash

# AIQuant 系统恢复脚本
# 用途: 从备份中恢复系统数据和配置
# 作者: AIQuant Team
# 版本: 1.0.0

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 配置变量
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
BACKUP_BASE_DIR="/opt/backups"
RESTORE_DIR="/tmp/aiquant_restore"

# 恢复选项
RESTORE_DATABASE=true
RESTORE_FILES=true
RESTORE_CONFIG=true
RESTORE_LOGS=false
FORCE_RESTORE=false
BACKUP_CURRENT=true
DRY_RUN=false

# 通知设置
ENABLE_NOTIFICATIONS=true
SLACK_WEBHOOK_URL=""
EMAIL_RECIPIENTS=""

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $(date '+%Y-%m-%d %H:%M:%S') $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $(date '+%Y-%m-%d %H:%M:%S') $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $(date '+%Y-%m-%d %H:%M:%S') $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $(date '+%Y-%m-%d %H:%M:%S') $1"
}

# 错误处理
error_exit() {
    log_error "$1"
    send_notification "error" "恢复失败: $1"
    exit 1
}

# 显示帮助信息
show_help() {
    cat << EOF
AIQuant 系统恢复脚本

用法: $0 [OPTIONS] BACKUP_PATH

参数:
    BACKUP_PATH             备份目录路径或备份文件路径

选项:
    -h, --help              显示帮助信息
    --no-db                 跳过数据库恢复
    --no-files              跳过文件恢复
    --no-config             跳过配置恢复
    --restore-logs          恢复日志文件
    --force                 强制恢复，不进行确认
    --no-backup             恢复前不备份当前数据
    --dry-run               模拟运行，不实际恢复
    --restore-dir DIR       指定恢复临时目录 (默认: $RESTORE_DIR)

云存储恢复:
    --from-s3 BUCKET/PATH   从AWS S3恢复
    --from-azure CONTAINER/PATH  从Azure Blob恢复
    --from-gcs BUCKET/PATH  从Google Cloud Storage恢复

示例:
    $0 /opt/backups/20240101_120000                    # 从本地备份恢复
    $0 --no-db /opt/backups/20240101_120000            # 恢复但跳过数据库
    $0 --from-s3 my-bucket/aiquant-backups/20240101   # 从S3恢复
    $0 --force --dry-run /opt/backups/latest           # 强制模拟恢复

EOF
}

# 解析命令行参数
parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_help
                exit 0
                ;;
            --no-db)
                RESTORE_DATABASE=false
                shift
                ;;
            --no-files)
                RESTORE_FILES=false
                shift
                ;;
            --no-config)
                RESTORE_CONFIG=false
                shift
                ;;
            --restore-logs)
                RESTORE_LOGS=true
                shift
                ;;
            --force)
                FORCE_RESTORE=true
                shift
                ;;
            --no-backup)
                BACKUP_CURRENT=false
                shift
                ;;
            --dry-run)
                DRY_RUN=true
                shift
                ;;
            --restore-dir)
                RESTORE_DIR="$2"
                shift 2
                ;;
            --from-s3)
                CLOUD_SOURCE="s3"
                CLOUD_PATH="$2"
                shift 2
                ;;
            --from-azure)
                CLOUD_SOURCE="azure"
                CLOUD_PATH="$2"
                shift 2
                ;;
            --from-gcs)
                CLOUD_SOURCE="gcs"
                CLOUD_PATH="$2"
                shift 2
                ;;
            -*)
                error_exit "未知选项: $1"
                ;;
            *)
                if [[ -z "$BACKUP_PATH" ]]; then
                    BACKUP_PATH="$1"
                fi
                shift
                ;;
        esac
    done
    
    # 检查必需参数
    if [[ -z "$BACKUP_PATH" && -z "$CLOUD_SOURCE" ]]; then
        error_exit "请指定备份路径或云存储源"
    fi
}

# 加载配置
load_config() {
    local config_file="$PROJECT_ROOT/.restore_config"
    if [[ -f "$config_file" ]]; then
        source "$config_file"
        log_info "已加载恢复配置文件"
    fi
    
    # 从环境变量加载配置
    SLACK_WEBHOOK_URL=${RESTORE_SLACK_WEBHOOK:-$SLACK_WEBHOOK_URL}
    EMAIL_RECIPIENTS=${RESTORE_EMAIL_RECIPIENTS:-$EMAIL_RECIPIENTS}
}

# 预检查
pre_restore_check() {
    log_info "执行恢复前检查..."
    
    # 检查必需的命令
    local required_commands=("docker" "docker-compose" "tar" "gzip")
    for cmd in "${required_commands[@]}"; do
        if ! command -v "$cmd" &> /dev/null; then
            error_exit "缺少必需命令: $cmd"
        fi
    done
    
    # 检查Docker服务
    if ! docker info &> /dev/null; then
        error_exit "Docker服务未运行"
    fi
    
    # 检查磁盘空间
    local available_space=$(df "$PROJECT_ROOT" 2>/dev/null | awk 'NR==2{print $4}' || echo "0")
    local required_space=5242880  # 5GB in KB
    if [[ $available_space -lt $required_space ]]; then
        error_exit "磁盘空间不足，需要至少5GB可用空间"
    fi
    
    # 检查备份路径
    if [[ -n "$BACKUP_PATH" && ! -d "$BACKUP_PATH" && ! -f "$BACKUP_PATH" ]]; then
        error_exit "备份路径不存在: $BACKUP_PATH"
    fi
    
    # 检查云存储工具
    if [[ -n "$CLOUD_SOURCE" ]]; then
        case "$CLOUD_SOURCE" in
            "s3")
                if ! command -v aws &> /dev/null; then
                    error_exit "AWS CLI未安装"
                fi
                ;;
            "azure")
                if ! command -v az &> /dev/null; then
                    error_exit "Azure CLI未安装"
                fi
                ;;
            "gcs")
                if ! command -v gsutil &> /dev/null; then
                    error_exit "Google Cloud SDK未安装"
                fi
                ;;
        esac
    fi
    
    log_success "预检查完成"
}

# 确认恢复操作
confirm_restore() {
    if [[ "$FORCE_RESTORE" == "true" || "$DRY_RUN" == "true" ]]; then
        return 0
    fi
    
    echo
    log_warning "警告: 此操作将恢复系统数据，可能会覆盖当前数据！"
    echo
    echo "恢复配置:"
    echo "  备份源: ${BACKUP_PATH:-$CLOUD_SOURCE:$CLOUD_PATH}"
    echo "  数据库恢复: $RESTORE_DATABASE"
    echo "  文件恢复: $RESTORE_FILES"
    echo "  配置恢复: $RESTORE_CONFIG"
    echo "  日志恢复: $RESTORE_LOGS"
    echo "  当前数据备份: $BACKUP_CURRENT"
    echo
    
    read -p "确认继续恢复操作? (yes/no): " -r
    if [[ ! $REPLY =~ ^[Yy][Ee][Ss]$ ]]; then
        log_info "恢复操作已取消"
        exit 0
    fi
}

# 从云存储下载备份
download_from_cloud() {
    if [[ -z "$CLOUD_SOURCE" ]]; then
        return 0
    fi
    
    log_info "从云存储下载备份..."
    
    local download_dir="$RESTORE_DIR/download"
    mkdir -p "$download_dir"
    
    if [[ "$DRY_RUN" == "true" ]]; then
        log_info "[DRY RUN] 将从 $CLOUD_SOURCE 下载: $CLOUD_PATH"
        return 0
    fi
    
    case "$CLOUD_SOURCE" in
        "s3")
            log_info "从AWS S3下载: $CLOUD_PATH"
            aws s3 sync "s3://$CLOUD_PATH" "$download_dir/" || error_exit "S3下载失败"
            ;;
        "azure")
            log_info "从Azure Blob下载: $CLOUD_PATH"
            az storage blob download-batch \
                --source "$CLOUD_PATH" \
                --destination "$download_dir" || error_exit "Azure下载失败"
            ;;
        "gcs")
            log_info "从Google Cloud Storage下载: $CLOUD_PATH"
            gsutil -m cp -r "gs://$CLOUD_PATH/*" "$download_dir/" || error_exit "GCS下载失败"
            ;;
    esac
    
    # 设置下载的备份路径
    BACKUP_PATH="$download_dir"
    log_success "云存储下载完成"
}

# 备份当前数据
backup_current_data() {
    if [[ "$BACKUP_CURRENT" != "true" ]]; then
        return 0
    fi
    
    log_info "备份当前数据..."
    
    local current_backup_dir="$RESTORE_DIR/current_backup_$(date +%Y%m%d_%H%M%S)"
    
    if [[ "$DRY_RUN" == "true" ]]; then
        log_info "[DRY RUN] 将备份当前数据到: $current_backup_dir"
        return 0
    fi
    
    mkdir -p "$current_backup_dir"
    
    cd "$PROJECT_ROOT"
    
    # 备份当前数据库
    if [[ "$RESTORE_DATABASE" == "true" ]]; then
        log_info "备份当前数据库..."
        if docker-compose exec -T postgres pg_isready -U aiquant; then
            docker-compose exec -T postgres pg_dump \
                -U aiquant -d aiquant \
                --format=custom \
                > "$current_backup_dir/current_postgres.dump" || log_warning "当前数据库备份失败"
        fi
    fi
    
    # 备份当前文件
    if [[ "$RESTORE_FILES" == "true" ]]; then
        log_info "备份当前文件..."
        tar -czf "$current_backup_dir/current_files.tar.gz" \
            data logs uploads models strategies backtest_results 2>/dev/null || log_warning "当前文件备份失败"
    fi
    
    # 备份当前配置
    if [[ "$RESTORE_CONFIG" == "true" ]]; then
        log_info "备份当前配置..."
        tar -czf "$current_backup_dir/current_config.tar.gz" \
            .env* docker-compose*.yml nginx ssl 2>/dev/null || log_warning "当前配置备份失败"
    fi
    
    log_success "当前数据备份完成: $current_backup_dir"
}

# 准备恢复环境
prepare_restore() {
    log_info "准备恢复环境..."
    
    if [[ "$DRY_RUN" == "true" ]]; then
        log_info "[DRY RUN] 将准备恢复环境"
        return 0
    fi
    
    # 创建恢复目录
    mkdir -p "$RESTORE_DIR"/{database,files,config,logs}
    
    # 停止服务
    log_info "停止AIQuant服务..."
    cd "$PROJECT_ROOT"
    docker-compose down || log_warning "服务停止失败"
    
    # 解压备份文件
    if [[ -f "$BACKUP_PATH" ]]; then
        log_info "解压备份文件..."
        tar -xzf "$BACKUP_PATH" -C "$RESTORE_DIR/" || error_exit "备份文件解压失败"
        
        # 查找解压后的备份目录
        local extracted_dir=$(find "$RESTORE_DIR" -maxdepth 1 -type d -name "20*" | head -1)
        if [[ -n "$extracted_dir" ]]; then
            BACKUP_PATH="$extracted_dir"
        fi
    fi
    
    log_success "恢复环境准备完成"
}

# 恢复数据库
restore_database() {
    if [[ "$RESTORE_DATABASE" != "true" ]]; then
        log_info "跳过数据库恢复"
        return 0
    fi
    
    log_info "开始恢复数据库..."
    
    if [[ "$DRY_RUN" == "true" ]]; then
        log_info "[DRY RUN] 将恢复数据库"
        return 0
    fi
    
    cd "$PROJECT_ROOT"
    
    # 启动数据库服务
    log_info "启动数据库服务..."
    docker-compose up -d postgres redis timescaledb || error_exit "数据库服务启动失败"
    
    # 等待数据库就绪
    log_info "等待数据库就绪..."
    local retry_count=0
    while ! docker-compose exec -T postgres pg_isready -U aiquant; do
        sleep 5
        ((retry_count++))
        if [[ $retry_count -gt 12 ]]; then
            error_exit "数据库启动超时"
        fi
    done
    
    # 恢复PostgreSQL
    local postgres_dump=$(find "$BACKUP_PATH" -name "*.dump" | head -1)
    if [[ -n "$postgres_dump" ]]; then
        log_info "恢复PostgreSQL数据库..."
        docker-compose exec -T postgres dropdb -U aiquant aiquant --if-exists || true
        docker-compose exec -T postgres createdb -U aiquant aiquant
        docker-compose exec -T postgres pg_restore \
            -U aiquant -d aiquant \
            --verbose --clean --if-exists \
            < "$postgres_dump" || log_warning "PostgreSQL恢复失败"
        log_success "PostgreSQL恢复完成"
    fi
    
    # 恢复TimescaleDB
    local tsdb_backup=$(find "$BACKUP_PATH" -name "timescaledb_*.sql.gz" | head -1)
    if [[ -n "$tsdb_backup" ]]; then
        log_info "恢复TimescaleDB..."
        zcat "$tsdb_backup" | docker-compose exec -T timescaledb psql \
            -U aiquant -d aiquant_ts || log_warning "TimescaleDB恢复失败"
        log_success "TimescaleDB恢复完成"
    fi
    
    # 恢复Redis
    local redis_backup=$(find "$BACKUP_PATH" -name "redis_*.rdb" | head -1)
    if [[ -n "$redis_backup" ]]; then
        log_info "恢复Redis数据..."
        docker-compose exec -T redis redis-cli FLUSHALL
        docker cp "$redis_backup" $(docker-compose ps -q redis):/data/dump.rdb
        docker-compose restart redis || log_warning "Redis恢复失败"
        log_success "Redis恢复完成"
    fi
}

# 恢复文件
restore_files() {
    if [[ "$RESTORE_FILES" != "true" ]]; then
        log_info "跳过文件恢复"
        return 0
    fi
    
    log_info "开始恢复应用文件..."
    
    if [[ "$DRY_RUN" == "true" ]]; then
        log_info "[DRY RUN] 将恢复应用文件"
        return 0
    fi
    
    cd "$PROJECT_ROOT"
    
    # 恢复应用文件
    local files_backup=$(find "$BACKUP_PATH" -name "application_files.tar.gz" | head -1)
    if [[ -n "$files_backup" ]]; then
        log_info "恢复应用文件..."
        tar -xzf "$files_backup" -C "$PROJECT_ROOT/" || log_warning "应用文件恢复失败"
        log_success "应用文件恢复完成"
    fi
    
    # 设置文件权限
    log_info "设置文件权限..."
    find data logs uploads models strategies backtest_results -type d -exec chmod 755 {} \; 2>/dev/null || true
    find data logs uploads models strategies backtest_results -type f -exec chmod 644 {} \; 2>/dev/null || true
}

# 恢复配置
restore_config() {
    if [[ "$RESTORE_CONFIG" != "true" ]]; then
        log_info "跳过配置恢复"
        return 0
    fi
    
    log_info "开始恢复配置文件..."
    
    if [[ "$DRY_RUN" == "true" ]]; then
        log_info "[DRY RUN] 将恢复配置文件"
        return 0
    fi
    
    cd "$PROJECT_ROOT"
    
    # 恢复应用配置
    local config_backup=$(find "$BACKUP_PATH" -name "configuration.tar.gz" | head -1)
    if [[ -n "$config_backup" ]]; then
        log_info "恢复应用配置..."
        tar -xzf "$config_backup" -C "$PROJECT_ROOT/" || log_warning "应用配置恢复失败"
        log_success "应用配置恢复完成"
    fi
    
    # 恢复系统配置
    local system_config_backup=$(find "$BACKUP_PATH" -name "system_config.tar.gz" | head -1)
    if [[ -n "$system_config_backup" ]]; then
        log_info "恢复系统配置..."
        tar -xzf "$system_config_backup" -C / || log_warning "系统配置恢复失败"
        log_success "系统配置恢复完成"
    fi
}

# 恢复日志
restore_logs() {
    if [[ "$RESTORE_LOGS" != "true" ]]; then
        log_info "跳过日志恢复"
        return 0
    fi
    
    log_info "开始恢复日志文件..."
    
    if [[ "$DRY_RUN" == "true" ]]; then
        log_info "[DRY RUN] 将恢复日志文件"
        return 0
    fi
    
    cd "$PROJECT_ROOT"
    
    # 恢复应用日志
    local logs_backup=$(find "$BACKUP_PATH" -name "application_logs.tar.gz" | head -1)
    if [[ -n "$logs_backup" ]]; then
        log_info "恢复应用日志..."
        tar -xzf "$logs_backup" -C "$PROJECT_ROOT/" || log_warning "应用日志恢复失败"
        log_success "应用日志恢复完成"
    fi
}

# 启动服务
start_services() {
    log_info "启动AIQuant服务..."
    
    if [[ "$DRY_RUN" == "true" ]]; then
        log_info "[DRY RUN] 将启动服务"
        return 0
    fi
    
    cd "$PROJECT_ROOT"
    
    # 启动所有服务
    docker-compose up -d || error_exit "服务启动失败"
    
    # 等待服务就绪
    log_info "等待服务就绪..."
    sleep 30
    
    # 检查服务状态
    local failed_services=()
    while read -r service status; do
        if [[ "$status" != "Up" ]]; then
            failed_services+=("$service")
        fi
    done < <(docker-compose ps --format "table {{.Service}}\t{{.Status}}" | tail -n +2)
    
    if [[ ${#failed_services[@]} -gt 0 ]]; then
        log_warning "以下服务启动失败: ${failed_services[*]}"
    else
        log_success "所有服务启动成功"
    fi
}

# 验证恢复
verify_restore() {
    log_info "验证恢复结果..."
    
    if [[ "$DRY_RUN" == "true" ]]; then
        log_info "[DRY RUN] 将验证恢复结果"
        return 0
    fi
    
    cd "$PROJECT_ROOT"
    
    local verification_failed=false
    
    # 验证数据库连接
    if [[ "$RESTORE_DATABASE" == "true" ]]; then
        if ! docker-compose exec -T postgres pg_isready -U aiquant; then
            log_error "PostgreSQL连接失败"
            verification_failed=true
        else
            log_success "PostgreSQL连接正常"
        fi
        
        if docker-compose ps redis | grep -q "Up"; then
            if ! docker-compose exec -T redis redis-cli ping | grep -q "PONG"; then
                log_error "Redis连接失败"
                verification_failed=true
            else
                log_success "Redis连接正常"
            fi
        fi
    fi
    
    # 验证关键文件
    if [[ "$RESTORE_FILES" == "true" ]]; then
        local key_dirs=("data" "models" "strategies")
        for dir in "${key_dirs[@]}"; do
            if [[ ! -d "$dir" ]]; then
                log_error "关键目录缺失: $dir"
                verification_failed=true
            fi
        done
    fi
    
    # 验证配置文件
    if [[ "$RESTORE_CONFIG" == "true" ]]; then
        local key_files=("docker-compose.yml" ".env.production")
        for file in "${key_files[@]}"; do
            if [[ ! -f "$file" ]]; then
                log_error "关键配置文件缺失: $file"
                verification_failed=true
            fi
        done
    fi
    
    if [[ "$verification_failed" == "true" ]]; then
        error_exit "恢复验证失败"
    else
        log_success "恢复验证通过"
    fi
}

# 清理恢复环境
cleanup_restore() {
    log_info "清理恢复环境..."
    
    if [[ "$DRY_RUN" == "true" ]]; then
        log_info "[DRY RUN] 将清理恢复环境"
        return 0
    fi
    
    # 保留当前数据备份，清理其他临时文件
    if [[ -d "$RESTORE_DIR" ]]; then
        find "$RESTORE_DIR" -name "current_backup_*" -type d | while read backup_dir; do
            log_info "保留当前数据备份: $backup_dir"
        done
        
        find "$RESTORE_DIR" -maxdepth 1 -type d ! -name "current_backup_*" ! -path "$RESTORE_DIR" -exec rm -rf {} \; 2>/dev/null || true
        find "$RESTORE_DIR" -maxdepth 1 -type f -exec rm -f {} \; 2>/dev/null || true
    fi
    
    log_success "恢复环境清理完成"
}

# 发送通知
send_notification() {
    if [[ "$ENABLE_NOTIFICATIONS" != "true" ]]; then
        return 0
    fi
    
    local status="$1"
    local message="$2"
    
    # Slack通知
    if [[ -n "$SLACK_WEBHOOK_URL" ]]; then
        local color="good"
        if [[ "$status" == "error" ]]; then
            color="danger"
        elif [[ "$status" == "warning" ]]; then
            color="warning"
        fi
        
        curl -X POST "$SLACK_WEBHOOK_URL" \
            -H 'Content-type: application/json' \
            --data "{
                \"attachments\": [{
                    \"color\": \"$color\",
                    \"title\": \"AIQuant系统恢复\",
                    \"text\": \"$message\",
                    \"fields\": [
                        {\"title\": \"备份源\", \"value\": \"${BACKUP_PATH:-$CLOUD_SOURCE:$CLOUD_PATH}\", \"short\": true},
                        {\"title\": \"恢复时间\", \"value\": \"$(date)\", \"short\": true}
                    ]
                }]
            }" || log_warning "Slack通知发送失败"
    fi
    
    # 邮件通知
    if [[ -n "$EMAIL_RECIPIENTS" ]]; then
        echo "$message" | mail -s "AIQuant恢复通知 - $status" "$EMAIL_RECIPIENTS" || log_warning "邮件通知发送失败"
    fi
}

# 主函数
main() {
    log_info "开始AIQuant系统恢复"
    log_info "备份源: ${BACKUP_PATH:-$CLOUD_SOURCE:$CLOUD_PATH}"
    log_info "时间: $(date)"
    
    # 解析命令行参数
    parse_args "$@"
    
    # 加载配置
    load_config
    
    # 执行预检查
    pre_restore_check
    
    # 确认恢复操作
    confirm_restore
    
    # 记录开始时间
    local start_time=$(date +%s)
    
    # 从云存储下载备份
    download_from_cloud
    
    # 备份当前数据
    backup_current_data
    
    # 准备恢复环境
    prepare_restore
    
    # 执行恢复
    restore_database
    restore_files
    restore_config
    restore_logs
    
    # 启动服务
    start_services
    
    # 验证恢复
    verify_restore
    
    # 清理恢复环境
    cleanup_restore
    
    # 计算恢复时间
    local end_time=$(date +%s)
    local duration=$((end_time - start_time))
    
    # 发送成功通知
    local success_message="AIQuant系统恢复成功完成
备份源: ${BACKUP_PATH:-$CLOUD_SOURCE:$CLOUD_PATH}
耗时: ${duration}秒
数据库恢复: $RESTORE_DATABASE
文件恢复: $RESTORE_FILES
配置恢复: $RESTORE_CONFIG"
    
    send_notification "success" "$success_message"
    
    log_success "恢复完成！"
    log_info "耗时: ${duration}秒"
    log_info "备份源: ${BACKUP_PATH:-$CLOUD_SOURCE:$CLOUD_PATH}"
    
    if [[ "$DRY_RUN" == "true" ]]; then
        log_info "这是模拟运行，没有实际执行恢复操作"
    fi
}

# 信号处理
cleanup_on_exit() {
    local exit_code=$?
    if [[ $exit_code -ne 0 ]]; then
        log_error "恢复过程中断，退出码: $exit_code"
        send_notification "error" "恢复过程异常中断，退出码: $exit_code"
    fi
    
    exit $exit_code
}

# 设置信号处理
trap cleanup_on_exit EXIT
trap 'error_exit "恢复被用户中断"' INT TERM

# 检查是否直接运行脚本
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi